package edu.usc.summerizer;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Hashtable;

import edu.usc.utils.FileUtils;
import edu.usc.utils.Parse;


public class EMTrial {

	public static String trainingURL = "C:\\Ramesh\\works\\projects\\java_ws\\StorySummerizer\\trainingsummary";
	public static String observedDataURL = "C:\\Ramesh\\works\\projects\\java_ws\\StorySummerizer\\resources\\observed.data";

	public static ArrayList<ArrayList<Double>> alphaTable = new ArrayList<>();
	public static ArrayList<ArrayList<Double>>  betaTable = new ArrayList<>();
	private static int NUMBER_OF_ITERATIONS = 25;
	private static String[] pocs = null; 

	public static void main(String[] args) {

		//testViterbi();
		// Populate the parse tables and assign them to the variables for EM 
		process();
	}

	public static void process() {
		Parse.populate(trainingURL);
		HashMap<String, Double> bigramProbs = new HashMap<>();
		HashMap<String, Double> biGramProbsBegin = new HashMap<>();
		HashMap<String, Double> biGramProbsEnd = new HashMap<>();
		HashMap<String, Double> wordPOCProbs = new HashMap<>();
		pocs = new String[Parse.POCProbs.keySet().size()];
		int index = 0;
		for (String poc: Parse.POCProbs.keySet()) {
			pocs[index++] = poc;
		}

		for (String bigram: Parse.POCPOCProbs.keySet()) {
			if(bigram.split("\\|")[1].equals("<s>")) {				
				biGramProbsBegin.put(bigram.split("\\|")[0], Parse.POCPOCProbs.get(bigram));
			} else if (bigram.split("\\|")[0].equals("</s>")) {
				biGramProbsEnd.put(bigram.split("\\|")[1], Parse.POCPOCProbs.get(bigram));
			} else {
				bigramProbs.put(bigram, Parse.POCPOCProbs.get(bigram));
			}
		}
		wordPOCProbs = Parse.wordPOCProbs;
		// Get the observed data. Stop, stem and put them into a list.
		ArrayList<String> observedData = new ArrayList<>();
		ArrayList<String> story = FileUtils.contentsOfTheFileAsList(observedDataURL);

		//for (String temp: story) {
		String temp = story.get(0);		
		temp  = FileUtils.cleanSentences(temp);
		/*String stopped = FileUtils.removeStopWordsString(temp, null);
			String stemmed = FileUtils.stemString(stopped);
			temp = stemmed;*/

		for(String w: temp.split(" ")) {
			if (w.length() > 0) {
				observedData.add(w);
			}
		}

		//System.out.println(observedData);
		//System.out.println(biGramProbsBegin);
		//System.out.println(wordPOCProbs.size());
		Hashtable<String, Double> counts = new Hashtable<>();
		while(NUMBER_OF_ITERATIONS-- != 0) {
			//System.out.print(NUMBER_OF_ITERATIONS);
			viterbiSearch(bigramProbs, biGramProbsBegin, 
					wordPOCProbs, pocs, observedData);
			createAlphas(biGramProbsBegin, biGramProbsEnd, 
					wordPOCProbs, observedData);
			createBetas(biGramProbsBegin, biGramProbsEnd, 
					wordPOCProbs, observedData);

			updateCounts(observedData, wordPOCProbs, counts);
			updateEstimations(observedData, wordPOCProbs, counts);
			alphaTable = new ArrayList<>();
			betaTable = new ArrayList<>();
		}
		System.out.println(temp);
		//}
	}

	public static void updateCounts(
			ArrayList<String> observedData,
			HashMap<String, Double> wordPOCProbs,
			Hashtable<String, Double> counts) {
		for (String s1: pocs) {
			for (String s2 : observedData) {
				counts.put(s1+","+s2, 0.0);
			}
		}
		double FINAL_ALPHA = alphaTable.get(alphaTable.size()-1).get(0);
		for(int i=0; i < pocs.length; i++) {
			int alphaIndex = 0;
			int betaIndex = betaTable.size()-3;
			for (int j=0; j < observedData.size(); j++) {
				double val = alphaTable.get(alphaIndex).get(i) + betaTable.get(betaIndex).get(i) - FINAL_ALPHA;
				if(wordPOCProbs.get(observedData.get(j)+"|"+pocs[i]) == null) {
					val = 0;
				} else {
					val = Math.exp(val) * wordPOCProbs.get(observedData.get(j)+"|"+pocs[i]);
				}
				double c = counts.get(pocs[i]+","+observedData.get(j));
				c+=val;
				counts.put(pocs[i]+","+observedData.get(j), c);
				alphaIndex += 2;
				betaIndex -= 2;
			}
		}
		//System.out.println(counts);
	}

	public static void updateEstimations(
			ArrayList<String> observedData,
			HashMap<String, Double> wordPOCProbs,
			Hashtable<String, Double> counts) {
		//encryptionProbs
		HashSet<String> uniqueEnciphers = new HashSet<>(observedData);
		Hashtable<String, Double> sums = new Hashtable<>();
		for(int i =0; i < pocs.length; i++) {
			double sum = 0.0;
			for(int j=0; j < uniqueEnciphers.size(); j++) {
				sum += counts.get(pocs[i]+","+observedData.get(j));
			}
			sums.put(pocs[i], sum);
		}

		for(int i =0; i < pocs.length; i++) {
			for(int j=0; j < observedData.size(); j++) {
				double count = counts.get(pocs[i]+","+observedData.get(j));
				count = count/sums.get(pocs[i]);
				wordPOCProbs.put(observedData.get(j)+"|"+pocs[i], count);
			}
		}
		//System.out.println(encryptionProbs);;
	}

	public static void testViterbi(){
		HashMap<String, Double> bigramProbsPOS = new HashMap<>();
		HashMap<String, Double> startBigramProbsPOS = new HashMap<>();
		HashMap<String, Double> wordPOSProbs = new HashMap<>();
		ArrayList<String> observedData = new ArrayList<>();
		String[] pocs = new String[] {"P", "V", "N", "A"};
		bigramProbsPOS.put("P|P", 0.0);
		bigramProbsPOS.put("P|V", 0.0);
		bigramProbsPOS.put("P|N", 0.0);
		bigramProbsPOS.put("P|A", 0.0);

		bigramProbsPOS.put("V|P", 0.6);
		bigramProbsPOS.put("V|V", 0.0);
		bigramProbsPOS.put("V|N", 0.2);
		bigramProbsPOS.put("V|A", 0.9);

		bigramProbsPOS.put("N|P", 0.0);
		bigramProbsPOS.put("N|V", 0.9);
		bigramProbsPOS.put("N|N", 0.7);
		bigramProbsPOS.put("N|A", 0.0);

		bigramProbsPOS.put("A|P", 0.3);
		bigramProbsPOS.put("A|V", 0.0);
		bigramProbsPOS.put("A|N", 0.0);
		bigramProbsPOS.put("A|A", 0.0);

		startBigramProbsPOS.put("P", 0.6);
		startBigramProbsPOS.put("V", 0.05);
		startBigramProbsPOS.put("N", 0.3);
		startBigramProbsPOS.put("A", 0.05);

		wordPOSProbs.put("T|P", 0.07);
		wordPOSProbs.put("T|V", 0.0);
		wordPOSProbs.put("T|N", 0.0);
		wordPOSProbs.put("T|A", 0.0);

		wordPOSProbs.put("C|P", 0.0);
		wordPOSProbs.put("C|V", 0.00001);
		wordPOSProbs.put("C|N", 0.0001);
		wordPOSProbs.put("C|A", 0.21);

		wordPOSProbs.put("F|P", 0.0);
		wordPOSProbs.put("F|V", 0.0001);
		wordPOSProbs.put("F|N", 0.0001);
		wordPOSProbs.put("F|A", 0.0);

		observedData.add("T");
		observedData.add("C");
		observedData.add("F");
		viterbiSearch(bigramProbsPOS, startBigramProbsPOS, wordPOSProbs, pocs, observedData);
	}

	public static String viterbiSearch(HashMap<String, Double> bigramProbs,
			HashMap<String, Double> biGramProbsBegin,
			HashMap<String, Double> wordPOCProbs,
			String[] pocs,
			ArrayList<String> observedData
			){
		ArrayList<ArrayList<Double>> weights= new ArrayList<>();
		ArrayList<ArrayList<Integer>> bestPreds = new ArrayList<>();
		ArrayList<Double> toAdd = new ArrayList<>();
		for(int i=0; i < pocs.length; i++) {
			if(biGramProbsBegin.get(pocs[i]) == null ||
					wordPOCProbs.get(observedData.get(0)+"|"+pocs[i]) == null) {
				toAdd.add(-99999.99);
			} else {
				toAdd.add(Math.log(biGramProbsBegin.get(pocs[i])) 
						+ Math.log(wordPOCProbs.get(observedData.get(0)+"|"+pocs[i])));
			}
		}
		weights.add(toAdd);

		for(int i=1; i < observedData.size(); i++) {
			toAdd = new ArrayList<>();
			ArrayList<Integer> toAddPred = new ArrayList<>();
			for(int j=0; j < pocs.length; j++) {
				double weightOfEntry = 0.0;
				int best_pred =0;
				double best_score = Double.NEGATIVE_INFINITY;
				for(int k=0; k < pocs.length; k++) {
					if ((bigramProbs.get(pocs[j]+"|"+pocs[k]) == null)
							|| wordPOCProbs.get(observedData.get(i)+"|"+pocs[j]) == null) {
						weightOfEntry = -99999.99;
					} else {
						//System.out.println(observedData.get(i)+"|"+pocs[j]);
						weightOfEntry = weights.get(i-1).get(k) +
								Math.log(bigramProbs.get(pocs[j]+"|"+pocs[k])) +
								Math.log(wordPOCProbs.get(observedData.get(i)+"|"+pocs[j]));
					}
					if(weightOfEntry >= best_score) {
						best_score = weightOfEntry;
						best_pred = k;						
					}
				}
				toAddPred.add(best_pred);								
				toAdd.add(best_score);
			}
			bestPreds.add(toAddPred);
			weights.add(toAdd);
		}

		Double bestWeight = Collections.max(weights.get(observedData.size()-1));
		int best_pred = weights.get(observedData.size()-1).indexOf(bestWeight);
		String solution = "";
		solution += pocs[best_pred];

		for(int i = observedData.size()-2; i >= 0; i--) {
			if(observedData.get(i).equals(" ")) {
				solution =  " " + solution;
			} else {
				best_pred = bestPreds.get(i).get(best_pred);
				solution= " " + solution;
				solution =  pocs[best_pred] + solution;
			}

		}
		//System.out.println(solution.length() +" ;;; "+encipheredMessage.size());
		System.out.println((25 - NUMBER_OF_ITERATIONS) + "\t" + bestWeight*-1);
		/*if(NUMBER_OF_ITERATIONS == 0) {
			System.out.print(solution + "\t");
			//System.out.println(solution + "\t" + bestWeight);
		}*/
		//FileUtils.writeToFileAlreadyExisting("decrypt", solution);
		return solution;
	}

	public static void createBetas(HashMap<String, Double> probsBegin,
			HashMap<String, Double> bigramProbs, // POC|POC
			HashMap<String, Double> wordPOCProbs,
			ArrayList<String> observedData
			){
		ArrayList<Double> presentBetas = new ArrayList<>();
		ArrayList<Double> previousBetas = new ArrayList<>();

		for(int i=0; i < pocs.length; i++) {
			presentBetas.add(0.0);
		}

		betaTable.add(presentBetas);
		previousBetas.addAll(presentBetas);
		presentBetas = new ArrayList<>();

		ArrayList<String> reverseEncipheredMessage = new ArrayList<>();
		for(int i=observedData.size()-1; i >= 0; i--) {
			reverseEncipheredMessage.add(observedData.get(i));
		}

		for (int enc = 0; enc < reverseEncipheredMessage.size(); enc++) {			
			String presentEncryption = reverseEncipheredMessage.get(enc);			
			for(int i=0 ; i < pocs.length; i++) {
				if(wordPOCProbs.get(presentEncryption+"|"+pocs[i]) 
						== null) {
					presentBetas.add(
							previousBetas.get(i));
				} else {
					presentBetas.add(
							previousBetas.get(i) + 
							Math.log(wordPOCProbs.get(presentEncryption+"|"+pocs[i]))
							);
				}
			}
			betaTable.add(presentBetas);

			previousBetas = new ArrayList<>();
			previousBetas.addAll(presentBetas);
			presentBetas = new ArrayList<>();
			if(enc ==  observedData.size() - 1) {
				double finalD = previousBetas.get(0);
				for(int i=1; i < previousBetas.size(); i++) {
					finalD = logAdd(finalD, previousBetas.get(i));
				}
				presentBetas.add(finalD);
			} else {
				for(int i=0 ; i < pocs.length; i++) {
					double newAlpha = 0.0;
					for(int j=0 ; j < pocs.length; j++) {
						if(bigramProbs.get(pocs[j]+"|"+pocs[i]) == null) {
						} else if (j==0) {
							newAlpha = previousBetas.get(j)+ Math.log(bigramProbs.get(pocs[j]+"|"+pocs[i]));
						} else {
							newAlpha = logAdd(newAlpha, previousBetas.get(j) + Math.log(bigramProbs.get(pocs[j]+"|"+pocs[i])));
						}
					}
					presentBetas.add(newAlpha);
				}
			}

			betaTable.add(presentBetas);
			previousBetas = new ArrayList<>();
			previousBetas.addAll(presentBetas);
			presentBetas = new ArrayList<>();
		}
		//System.out.println(betaTable);
	}

	public static void createAlphas(HashMap<String, Double> probsBegin,
			HashMap<String, Double> bigramProbs, // POC|POC
			HashMap<String, Double> wordPOCProbs,
			ArrayList<String> observedData
			){
		ArrayList<Double> presentAlphas = new ArrayList<>();
		ArrayList<Double> previousAlphas = new ArrayList<>();

		for(int i=0; i < pocs.length; i++) {
			//System.out.println(ALPHAS[i] + " :: "+unigramProbs);
			if (probsBegin.get(pocs[i]) == null) {
				presentAlphas.add(-9999.9);
			} else {
				presentAlphas.add(Math.log(probsBegin.get(pocs[i])));
			}
		}

		alphaTable.add(presentAlphas);
		previousAlphas.addAll(presentAlphas);
		presentAlphas = new ArrayList<>();

		for (int enc = 0; enc < observedData.size(); enc++) {			
			String presentEncryption = observedData.get(enc);			
			for(int i=0 ; i < pocs.length; i++) {
				//System.out.println(presentEncryption+"|"+ALPHAS[i]);
				if(wordPOCProbs.get(presentEncryption+"|"+pocs[i]) 
						== null) {
					presentAlphas.add(previousAlphas.get(i));
				} else {
					presentAlphas.add(
							previousAlphas.get(i) + 
							Math.log(wordPOCProbs.get(presentEncryption+"|"+pocs[i])));
				}
			}
			alphaTable.add(presentAlphas);
			previousAlphas = new ArrayList<>();			
			previousAlphas.addAll(presentAlphas);
			presentAlphas = new ArrayList<>();

			if(enc ==  observedData.size() - 1) {
				double finalD = previousAlphas.get(0);
				for(int i=1; i < previousAlphas.size(); i++) {
					finalD = logAdd(finalD, previousAlphas.get(i));
				}
				presentAlphas.add(finalD);
			} else {

				for(int i=0 ; i < pocs.length; i++) {
					double newAlpha = 0.0;
					for(int j=0 ; j < pocs.length; j++) {
						if(bigramProbs.get(pocs[i]+"|"+pocs[j]) == null) {
						} else if (j==0) {
							newAlpha = previousAlphas.get(j)+ Math.log(bigramProbs.get(pocs[i]+"|"+pocs[j]));
						} else {
							newAlpha = logAdd(newAlpha, previousAlphas.get(j)+ Math.log(bigramProbs.get(pocs[i]+"|"+pocs[j])));
						}
					}
					presentAlphas.add(newAlpha);
				}
			}
			alphaTable.add(presentAlphas);
			previousAlphas = new ArrayList<>();
			previousAlphas.addAll(presentAlphas);
			presentAlphas = new ArrayList<>();
		}
		//System.out.println(alphaTable);
	}

	public static double logAdd(double logx, double logy) {
		if (logx == Double.NEGATIVE_INFINITY) {
			return logy;
		}

		if(logy == Double.NEGATIVE_INFINITY) {
			return logx;
		}

		if( (logx-logy) > 16) {
			return logx;
		}

		if( (logy-logx) > 16) {
			return logy;
		}

		if(logx > logy) {
			return (logx + Math.log(1.0 + Math.exp(logy-logx)));
		}

		return (logy + Math.log(1.0 + Math.exp(logx-logy)));
	}

}
